#!/usr/bin/env python
# -*- coding: utf-8 -*-


import os
import sys
import glob
import h5py
import numpy as np
from scipy.spatial.transform import Rotation
from torch.utils.data import Dataset


# Part of the code is referred from: https://github.com/charlesq34/pointnet

def download():
	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
	DATA_DIR = os.path.join(BASE_DIR, 'data')
	if not os.path.exists(DATA_DIR):
		os.mkdir(DATA_DIR)
	if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
		www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip'
		zipfile = os.path.basename(www)
		os.system('wget %s; unzip %s' % (www, zipfile))
		os.system('mv %s %s' % (zipfile[:-4], DATA_DIR))
		os.system('rm %s' % (zipfile))

def download_data(file):
	print_('################### Downloading Data ###################', color='g', style='bold')
	from google_drive_downloader import GoogleDriveDownloader as gdd

	if file=='train_data':
		file_id = '16YU-tdayVNBwM3XlPDgFrrzlPjhQN3PB'
	elif file=='car_data':
		file_id = '1k9W75uhUFTfA_iK7YePGn5t9f4JhtgSe'

	if not os.path.exists(os.path.join(os.getcwd(),'data',file)):
		gdd.download_file_from_google_drive(file_id=file_id,
										dest_path=os.path.join(os.getcwd(),'data',file+'.zip'),
										showsize=True,
										unzip=True)

		os.remove(os.path.join(os.getcwd(),'data',file+'.zip'))
	return True


def load_data(partition):
	download()
	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
	DATA_DIR = os.path.join(BASE_DIR, 'data')
	all_data = []
	all_label = []
	for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5' % partition)):
		f = h5py.File(h5_name)
		data = f['data'][:].astype('float32')
		label = f['label'][:].astype('int64')
		f.close()
		all_data.append(data)
		all_label.append(label)
	all_data = np.concatenate(all_data, axis=0)
	all_label = np.concatenate(all_label, axis=0)
	return all_data, all_label

def load_pcr_single(partition, idx=0):
	train_len = 5070
	test_len = 1024
	f = h5py.File('./data/train_data/templates.h5')
	data = f['templates'][idx].astype('float32')[:1024]
	f.close()
	if partition == 'train':
		new_data = np.tile(data, [train_len, 1]).reshape(-1, 1024, 3)
	if partition == 'test':
		new_data = np.tile(data, [test_len, 1]).reshape(-1, 1024, 3)
	return new_data

def load_pcr(partition):
	f = h5py.File('./data/train_data/templates.h5')
	data = f['templates'][:].astype('float32')
	f.close()
	if partition == 'test':
		data = data[:1024]
	return data

def read_poses(filename):
    import csv
    with open(os.path.join('data/train_data',filename),'r') as csvfile:
        csvreader = csv.reader(csvfile)
        poses = []
        for row in csvreader:
            row = [float(i) for i in row]
            poses.append(row)
    return np.asarray(poses)


def translate_pointcloud(pointcloud):
	xyz1 = np.random.uniform(low=2. / 3., high=3. / 2., size=[3])
	xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3])

	translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32')
	return translated_pointcloud


def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.05):
	N, C = pointcloud.shape
	pointcloud += np.clip(sigma * np.random.randn(N, C), -1 * clip, clip)
	return pointcloud


class ModelNet40(Dataset):
	def __init__(self, num_points, partition='train', gaussian_noise=False, unseen=False, factor=4):
		self.data, self.label = load_data(partition)
		self.num_points = num_points
		self.partition = partition
		self.gaussian_noise = gaussian_noise
		self.unseen = unseen
		self.label = self.label.squeeze()
		self.factor = factor
		if self.unseen:
			######## simulate testing on first 20 categories while training on last 20 categories
			if self.partition == 'test':
				self.data = self.data[self.label>=20]
				self.label = self.label[self.label>=20]
			elif self.partition == 'train':
				self.data = self.data[self.label<20]
				self.label = self.label[self.label<20]

	def __getitem__(self, item):
		pointcloud = self.data[item][:self.num_points]
		if self.gaussian_noise:
			pointcloud = jitter_pointcloud(pointcloud)
		if self.partition != 'train':
			np.random.seed(item)
		anglex = np.random.uniform() * np.pi / self.factor
		angley = np.random.uniform() * np.pi / self.factor
		anglez = np.random.uniform() * np.pi / self.factor

		cosx = np.cos(anglex)
		cosy = np.cos(angley)
		cosz = np.cos(anglez)
		sinx = np.sin(anglex)
		siny = np.sin(angley)
		sinz = np.sin(anglez)
		Rx = np.array([[1, 0, 0],
						[0, cosx, -sinx],
						[0, sinx, cosx]])
		Ry = np.array([[cosy, 0, siny],
						[0, 1, 0],
						[-siny, 0, cosy]])
		Rz = np.array([[cosz, -sinz, 0],
						[sinz, cosz, 0],
						[0, 0, 1]])
		R_ab = Rx.dot(Ry).dot(Rz)
		R_ba = R_ab.T
		#translation_ab = np.array([np.random.uniform(-0.5, 0.5), np.random.uniform(-0.5, 0.5),
		#                           np.random.uniform(-0.5, 0.5)])
		translation_ab = np.array([0.0, 0.0, 0.0])
		translation_ba = -R_ba.dot(translation_ab)

		pointcloud1 = pointcloud.T

		rotation_ab = Rotation.from_euler('zyx', [anglez, angley, anglex])
		pointcloud2 = rotation_ab.apply(pointcloud1.T).T + np.expand_dims(translation_ab, axis=1)

		euler_ab = np.asarray([anglez, angley, anglex])
		euler_ba = -euler_ab[::-1]

		pointcloud1 = np.random.permutation(pointcloud1.T).T
		pointcloud2 = np.random.permutation(pointcloud2.T).T

		return pointcloud1.astype('float32'), pointcloud2.astype('float32'), R_ab.astype('float32'), \
			   translation_ab.astype('float32'), R_ba.astype('float32'), translation_ba.astype('float32'), \
			   euler_ab.astype('float32'), euler_ba.astype('float32')

	def __len__(self):
		return self.data.shape[0]

class pcr(Dataset):
	def __init__(self, num_points, partition='train', gaussian_noise=True, unseen=False, factor=4):
		assert download_data('train_data'), "Error in downloading the dataset!!"
		self.data = load_pcr(partition)
		if partition == 'train':
			self.poses = read_poses('itr_net_train_data45.csv')[:len(self.data)]
		elif partition == 'test':
			self.poses = read_poses('itr_net_eval_data45.csv')[:len(self.data)]
		self.num_points = num_points
		self.partition = partition
		self.gaussian_noise = gaussian_noise
		self.unseen = unseen
		self.factor = factor

	def __getitem__(self, item):
		pointcloud = self.data[item][:self.num_points]
		trans_x, trans_y, trans_z, anglex, angley, anglez = self.poses[item]
		if self.partition != 'train':
			np.random.seed(item)

		cosx = np.cos(anglex)
		cosy = np.cos(angley)
		cosz = np.cos(anglez)
		sinx = np.sin(anglex)
		siny = np.sin(angley)
		sinz = np.sin(anglez)
		Rx = np.array([[1, 0, 0],
						[0, cosx, -sinx],
						[0, sinx, cosx]])
		Ry = np.array([[cosy, 0, siny],
						[0, 1, 0],
						[-siny, 0, cosy]])
		Rz = np.array([[cosz, -sinz, 0],
						[sinz, cosz, 0],
						[0, 0, 1]])
		R_ab = Rx.dot(Ry).dot(Rz)
		R_ba = R_ab.T
		translation_ab = np.array([0.0, 0.0, 0.0])
		translation_ba = -R_ba.T.dot(translation_ab)

		pointcloud1 = pointcloud.T

		rotation_ab = Rotation.from_euler('zyx', [anglez, angley, anglex])
		pointcloud2 = rotation_ab.apply(pointcloud1.T).T + np.expand_dims(translation_ab, axis=1)

		if self.gaussian_noise:
			pointcloud2 = add_noise(pointcloud2)

		euler_ab = np.asarray([anglez, angley, anglex])
		euler_ba = -euler_ab[::-1]

		# pointcloud1 = np.random.permutation(pointcloud1.T).T
		# pointcloud2 = np.random.permutation(pointcloud2.T).T

		return pointcloud1.astype('float32'), pointcloud2.astype('float32'), R_ab.astype('float32'), \
			   translation_ab.astype('float32'), R_ba.astype('float32'), translation_ba.astype('float32'), \
			   euler_ab.astype('float32'), euler_ba.astype('float32')

	def __len__(self):
		return self.data.shape[0]


class pcr_single(Dataset):
	def __init__(self, num_points, partition='train', gaussian_noise=False, unseen=False, factor=4):
		assert download_data('car_data'), "Error in downloading the dataset!!"
		self.data = load_pcr_single(partition, idx=2)
		if partition == 'train':
			self.poses = read_poses('itr_net_train_data45.csv')[:len(self.data)]
		elif partition == 'test':
			self.poses = read_poses('itr_net_eval_data45.csv')[:len(self.data)]
		self.num_points = num_points
		self.partition = partition
		self.gaussian_noise = gaussian_noise
		self.unseen = unseen
		self.factor = factor

	def __getitem__(self, item):
		pointcloud = self.data[item][:self.num_points]
		trans_x, trans_y, trans_z, anglex, angley, anglez = self.poses[item]
		if self.partition != 'train':
			np.random.seed(item)

		cosx = np.cos(anglex)
		cosy = np.cos(angley)
		cosz = np.cos(anglez)
		sinx = np.sin(anglex)
		siny = np.sin(angley)
		sinz = np.sin(anglez)
		Rx = np.array([[1, 0, 0],
						[0, cosx, -sinx],
						[0, sinx, cosx]])
		Ry = np.array([[cosy, 0, siny],
						[0, 1, 0],
						[-siny, 0, cosy]])
		Rz = np.array([[cosz, -sinz, 0],
						[sinz, cosz, 0],
						[0, 0, 1]])
		R_ab = Rx.dot(Ry).dot(Rz)
		R_ba = R_ab.T
		translation_ab = np.array([trans_x, trans_y, trans_z])
		translation_ba = -R_ba.T.dot(translation_ab)

		pointcloud1 = pointcloud.T

		rotation_ab = Rotation.from_euler('zyx', [anglez, angley, anglex])
		pointcloud2 = rotation_ab.apply(pointcloud1.T).T + np.expand_dims(translation_ab, axis=1)

		if self.gaussian_noise:
			pointcloud2 = add_noise(pointcloud2)

		euler_ab = np.asarray([anglez, angley, anglex])
		euler_ba = -euler_ab[::-1]

		# pointcloud1 = np.random.permutation(pointcloud1.T).T
		# pointcloud2 = np.random.permutation(pointcloud2.T).T

		return pointcloud1.astype('float32'), pointcloud2.astype('float32'), R_ab.astype('float32'), \
			   translation_ab.astype('float32'), R_ba.astype('float32'), translation_ba.astype('float32'), \
			   euler_ab.astype('float32'), euler_ba.astype('float32')

	def __len__(self):
		return self.data.shape[0]


if __name__ == '__main__':
	train = ModelNet40(1024)
	test = ModelNet40(1024, 'test')
	for data in train:
		print(len(data))
		break
